
# load custom packages from src dir
import sys
sys.path.insert(0, '..')

# python packages
import logging
import os
import time
import numpy as np
import pandas as pd

# Viz
import matplotlib.pyplot as plt
%matplotlib inline

# custom packages
import src.commons.dataset as ds
import src.commons.constants as cons

# U-Net custom implementation
from src.models.u_net import *

# Define the logging level
logging.getLogger().setLevel(logging.INFO)


# device
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device: ', DEVICE)

# hyperparameters
NUM_CLASSES = 3 # in anomaly detection, it is just the 3 channels
NUM_WORKERS = 2
NUM_EPOCHS = 5
BATCH_SIZE = 4
LEARNING_RATE = 1e-5
RESIZE_DIM = 100
CHECKPOINT_PATH = ".." + cons.DIR_SEP + os.path.join("results")
print("CHECKPOINT PATH: ", CHECKPOINT_PATH)


CATEGORY = "screw"
data = load_training_and_validation_data(CATEGORY, batch_size=BATCH_SIZE)
#train_dataset, train_dataloader, val_dataloader = data['dataset'], data['train'], data['val']
get_plot_training_data(data['dataset'], 4)


%%time
# Training
train_loss, val_loss, model = train_model(data, CATEGORY, 
    num_classes=NUM_CLASSES,
    device=DEVICE,
    resize_dim=RESIZE_DIM,
    learning_rate=LEARNING_RATE,
    checkpoint_path=CHECKPOINT_PATH,
    num_epochs=NUM_EPOCHS
)


print_loss(train_loss, val_loss)


test_data = load_test_data(CATEGORY, batch_size=BATCH_SIZE)
test_data, test_dataloader = test_data['dataset'], test_data['test']


%%time
run_predictions_and_plot_results(model, test_dataloader, device=DEVICE, resize_dim=RESIZE_DIM)


all_categories = [
    'bottle', 'cable', 'capsule', 
    'carpet', 'grid', 'hazelnut', 
    'leather', 'metal_nut', 'pill', 
    'screw', 'tile', 'toothbrush', 
    'transistor', 'wood', 'zipper']

explore_dimensions = [100, 200, 400]


# device
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device: ', DEVICE)

# hyperparameters
NUM_CLASSES = 3 # in anomaly detection, it is just the 3 channels
NUM_WORKERS = 2
NUM_EPOCHS = 20
BATCH_SIZE = 4
LEARNING_RATE = 1e-5
RESIZE_DIM = 100
CHECKPOINT_PATH = ".." + cons.DIR_SEP + os.path.join("results")
print("CHECKPOINT PATH: ", CHECKPOINT_PATH)

device:  cpu
CHECKPOINT PATH:  ../results


def run_for_category(cat, explore_dimensions):
    data = load_training_and_validation_data(cat, batch_size=BATCH_SIZE)
    get_plot_training_data(data['dataset'], 4)

    test_data = load_test_data(cat, batch_size=BATCH_SIZE)
    test_data, test_dataloader = test_data['dataset'], test_data['test']

    experiment, performance_train, performance_val, process_duration = [], [], [], []
    for resize_dim in explore_dimensions:
        ts_start = time.time()
        # Training
        train_loss, val_loss, model = train_model(data, cat, 
            num_classes=NUM_CLASSES,
            device=DEVICE,
            resize_dim=resize_dim,
            learning_rate=LEARNING_RATE,
            checkpoint_path=CHECKPOINT_PATH,
            num_epochs=NUM_EPOCHS
        )
    
        performance_train.append(np.min(train_loss))
        performance_val.append(np.min(val_loss))
        experiment.append(f"cat_{cat}_with_{resize_dim}_resize_dim")
    
        print_loss(train_loss, val_loss)
    
        run_predictions_and_plot_results(model, test_dataloader, device=DEVICE, resize_dim=resize_dim)
        end_ts = time.time()
        process_duration = end_ts - ts_start

    return pd.DataFrame({
        'experiment': experiment, 
        'best_train_perf': performance_train, 
        'best_val_perf': performance_val,
        'exp_duration_sec': process_duration
    })


def run_all_categories(all_categories, all_reshaping_sizes):
    df_results = pd.DataFrame(columns=['experiment', 'best_train_perf', 'best_val_perf', 'exp_duration_sec'])
    for cat in all_categories:
        df_cat_results = run_for_category(cat, all_reshaping_sizes)
        df_results = pd.concat([df_results, df_cat_results])
        # checkpoint the experiments results
        df_results.to_csv(os.path.join(CHECKPOINT_PATH, f"training_results_autoencoder_{cat}.csv"), index=False)
    return df_results


%%time
df_results = run_all_categories(all_categories, explore_dimensions)

209 209

Total trainable parameters: 3.10e+07

Epoch 1/20

loss = 0.197: 100%|██████████| 48/48 [00:34<00:00,  1.37it/s]

Train loss: 0.256

loss = 0.240: 100%|██████████| 5/5 [00:01<00:00,  3.95it/s]

Validation loss: 0.239

Epoch 2/20

loss = 0.137: 100%|██████████| 48/48 [00:35<00:00,  1.36it/s]

Train loss: 0.163

loss = 0.127: 100%|██████████| 5/5 [00:01<00:00,  4.10it/s]

Validation loss: 0.127

Epoch 3/20

CNN Autoencoder for anomaly detection¶

Model Parameters¶

Data¶

Train the model¶

Test¶

Run sequence¶